"""
LightRAG基础使用示例
"""

import os
import asyncio
from pathlib import Path
import logging

from lightrag import LightRAG

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

async def main():
    # 检查环境变量
    if not os.getenv("OPENAI_API_KEY"):
        raise ValueError("请设置OPENAI_API_KEY环境变量")
        
    # 创建LightRAG实例
    rag = LightRAG()
    
    # 添加示例文档
    docs_dir = Path(__file__).parent / "docs"
    if not docs_dir.exists():
        docs_dir.mkdir(parents=True)
        
    # 创建示例文档
    sample_text = """
    # 检索增强生成（RAG）简介
    
    检索增强生成（Retrieval-Augmented Generation，RAG）是一种将检索系统与生成模型相结合的技术。
    它通过检索相关文档来增强语言模型的知识，从而生成更准确、更可靠的回答。
    
    ## RAG的主要优势
    
    1. 知识更新：RAG可以通过更新检索库来更新知识，无需重新训练模型
    2. 可解释性：可以展示用于生成答案的参考文档
    3. 准确性：基于实际文档生成答案，减少幻觉
    
    ## RAG的工作流程
    
    1. 文档处理：将文档分割成小块并建立索引
    2. 检索：根据用户查询检索相关文档
    3. 生成：将检索到的文档作为上下文生成答案
    
    ## 应用场景
    
    - 问答系统
    - 客服机器人
    - 文档摘要
    - 知识库查询
    """
    
    sample_doc_path = docs_dir / "rag_intro.md"
    with open(sample_doc_path, "w", encoding="utf-8") as f:
        f.write(sample_text)
        
    # 添加文档
    await rag.add_document(sample_doc_path)
    
    # 测试查询
    questions = [
        "什么是RAG？",
        "RAG有什么优势？",
        "RAG的工作流程是怎样的？",
        "RAG可以应用在哪些场景？"
    ]
    
    for question in questions:
        print(f"\n问题: {question}")
        print("-" * 50)
        
        # 执行查询
        response = await rag.query(question)
        
        # 输出结果
        print(f"\n回答: {response.answer}")
        print(f"\n置信度: {response.score:.2f}")
        
        print("\n参考来源:")
        for source in response.sources:
            chunk = source["chunk"]
            print(f"- {chunk.metadata['file_name']} (相关度: {source['score']:.2f})")
            
        if response.followup_questions:
            print("\n相关问题:")
            for q in response.followup_questions:
                print(f"- {q}")
                
        print("=" * 50)
        
if __name__ == "__main__":
    asyncio.run(main()) 